package pers.whj.util.tree;

import pers.whj.util.list.QueueList;
import pers.whj.util.list.Stack;

public class AVLTreeImpl<T extends Comparable<T>> extends BinarySearchTreeImpl<T> implements AVLTree<T> {
    public AVLTreeImpl() {
    }

    protected static class Node<T extends Comparable<T>> extends AbstractTree.Node<T> {
        // the balance factor of this node, must be -1, 0 or 1
        protected int bf;

        public Node(T data) {
            super(data);
            this.bf = 0;
        }

        public Node(T data, Node<T> left, Node<T> right) {
            super(data, left, right);
            if (left == null && right == null) {
                this.bf = 0;
            } else if (left == null) {
                this.bf = 1;
            } else {
                this.bf = -1;
            }
        }

        @Override
        public void reset() {
            super.reset();
            this.bf = 0;
        }
    }

    @Override
    public int getBalanceFactor() {
        return this.root == null ? 0 : ((Node<T>) this.root).bf;
    }

    @Override
    public int getBalanceFactor(T data) {
        Node<T> node = (Node<T>) search(this.root, data);
        checkNode(node);
        return node.bf;
    }

    @Override
    public boolean insert(T data) {
        checkInsertData(data);
        if (this.root == null) {
            this.root = new Node<>(data);
            this.size = 1;
            return true;
        }
        Node<T> cur = (Node<T>) this.root;
        Node<T> pre = null;
        Stack<Node<T>> stack = new QueueList<>();
        while (cur != null) {
            if (cur.data.compareTo(data) == 0) {
                return false;
            }
            pre = cur;
            stack.push(pre);
            if (cur.data.compareTo(data) < 0) {
                cur = (Node<T>) cur.right;
            } else {
                cur = (Node<T>) cur.left;
            }
        }
        cur = new Node<>(data);
        if (pre.data.compareTo(data) < 0) {
            pre.right = cur;
        } else {
            pre.left = cur;
        }
        while (!stack.isEmpty()) {
            pre = stack.pop();
            if (pre.left == cur) {
                pre.bf--;
            } else {
                pre.bf++;
            }
            if (pre.bf == 0) {
                cur = pre;
                break;
            }
            if (pre.bf == 1 || pre.bf == -1) {
                cur = pre;
            } else {
                int flag = (pre.bf < 0) ? -1 : 1;
                if (cur.bf == flag) {
                    if (flag == -1) {
                        cur = rotateR(pre);
                    } else {
                        cur = rotateL(pre);
                    }
                } else {
                    if (flag == -1) {
                        cur = rotateLR(pre);
                    } else {
                        cur = rotateRL(pre);
                    }
                }
                break;
            }
        }
        if (stack.isEmpty()) {
            this.root = cur;
        } else {
            pre = stack.pop();
            if (pre.data.compareTo(cur.data) < 0) {
                pre.right = cur;
            } else {
                pre.left = cur;
            }
        }
        this.size++;
        return true;
    }

    @Override
    public boolean remove(T data) {
        if (data == null) {
            return false;
        }
        Node<T> p = (Node<T>) this.root;
        Node<T> pp = null;
        Node<T> ppp = null;
        Stack<Node<T>> stack = new QueueList<>();
        // find the position to delete
        while (p != null) {
            if (p.data.compareTo(data) == 0) {
                // find the data
                break;
            }
            pp = p;
            stack.push(pp);
            if (p.data.compareTo(data) < 0) {
                p = (Node<T>) p.right;
            } else {
                p = (Node<T>) p.left;
            }
        }
        if (p == null) {
            // fail to remove
            return false;
        }
        // has two children
        Node<T> q;
        if (p.left != null && p.right != null) {
            pp = p;
            stack.push(pp);
            q = (Node<T>) p.left;
            while (q.right != null) {
                pp = q;
                stack.push(pp);
                q = (Node<T>) q.right;
            }
            p.data = q.data;
            // to delete the node q
            p = q;
        }
        if (p.left != null) {
            q = (Node<T>) p.left;
        } else {
            q = (Node<T>) p.right;
        }
        // delete the root node
        if (pp == null) {
            this.root.reset();
            this.root = q;
            this.size--;
            return true;
        }
        // delete p and link nodes
        if (pp.left == p) {
            pp.left = q;
        } else {
            pp.right = q;
        }
        // check bf
        int d = 0, dd = 0;
        while (!stack.isEmpty()) {
            pp = stack.pop();
            if (pp.right == q) {
                pp.bf--;
            } else {
                pp.bf++;
            }
            if (!stack.isEmpty()) {
                ppp = stack.peek();
                dd = (ppp.left == pp) ? -1 : 1;
            } else {
                dd = 0;
            }
            if (pp.bf == 1 || pp.bf == -1) {
                break;
            }
            if (pp.bf != 0) {
                if (pp.bf < 0) {
                    d = -1;
                    q = (Node<T>) pp.left;
                } else {
                    d = 1;
                    q = (Node<T>) pp.right;
                }
                if (q.bf == 0) {
                    if (d == -1) {
                        pp = rotateR(pp);
                        pp.bf = 1;
                        ((Node<T>) pp.right).bf = -1;
                    } else {
                        pp = rotateL(pp);
                        pp.bf = -1;
                        ((Node<T>) pp.left).bf = 1;
                    }
                    break;
                }
                if (q.bf == d) {
                    if (d == -1) {
                        pp = rotateR(pp);
                    } else {
                        pp = rotateL(pp);
                    }
                } else {
                    if (d == -1) {
                        pp = rotateLR(pp);
                    } else {
                        pp = rotateRL(pp);
                    }
                }
                if (dd == -1) {
                    ppp.left = pp;
                } else if (dd == 1) {
                    ppp.right = pp;
                }
            }
            q = pp;
        }
        if (stack.isEmpty()) {
            this.root = pp;
        }
        this.size--;
        return true;
    }

    /**
     * single left rotation
     *
     * @param node from current node
     * @return the root node of this subtree after rotation
     */
    protected Node<T> rotateL(Node<T> node) {
        Node<T> subR = (Node<T>) node.right;
        node.right = subR.left;
        subR.left = node;
        node.bf = 0;
        subR.bf = 0;
        return subR;
    }

    protected Node<T> rotateR(Node<T> node) {
        Node<T> subL = (Node<T>) node.left;
        node.left = subL.right;
        subL.right = node;
        node.bf = 0;
        subL.bf = 0;
        return subL;
    }

    protected Node<T> rotateLR(Node<T> node) {
        Node<T> subL = (Node<T>) node.left;
        Node<T> root = (Node<T>) subL.right;
        subL.right = root.left;
        root.left = subL;
        if (root.bf <= 0) {
            subL.bf = 0;
        } else {
            subL.bf = -1;
        }
        node.left = root.right;
        root.right = node;
        if (root.bf == -1) {
            node.bf = 1;
        } else {
            node.bf = 0;
        }
        root.bf = 0;
        return root;
    }

    protected Node<T> rotateRL(Node<T> node) {
        Node<T> subR = (Node<T>) node.right;
        Node<T> root = (Node<T>) subR.left;
        subR.left = root.right;
        root.right = subR;
        if (root.bf >= 0) {
            subR.bf = 0;
        } else {
            subR.bf = 1;
        }
        node.right = root.left;
        root.left = node;
        if (root.bf == 1) {
            node.bf = -1;
        } else {
            node.bf = 0;
        }
        root.bf = 0;
        return root;
    }

    @Override
    protected boolean remove(T data, AbstractTree.Node<T> node, AbstractTree.Node<T> parent) {
        return remove(data);
    }
}